import gensim
import gensim.downloader as api
from gensim.models import KeyedVectors
from gensim.scripts.glove2word2vec import glove2word2vec
from nltk.corpus import movie_reviews
import random
from nltk import word_tokenize
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.svm import LinearSVC
from sklearn.decomposition import PCA

random.seed(42)
def load_movie_reviews():
    pos_ids = movie_reviews.fileids('pos')
    neg_ids = movie_reviews.fileids('neg')

    all_reviews = []
    for pids in pos_ids:
        all_reviews.append((movie_reviews.raw(pids), 'positive'))

    for nids in neg_ids:
        all_reviews.append((movie_reviews.raw(nids), 'negative'))

    random.shuffle(all_reviews)
    train_reviews = all_reviews[:1600]
    test_reviews = all_reviews[1600:]

    return train_reviews, test_reviews

# 将文本中每个词的词向量的平均作为文本的表示
def convert_text_to_vector(text, vectors):
    vector = []
    for word in word_tokenize(text):
        if word in vectors:
            vector.append(vectors[word])
    x=np.array(vector).T
    X_scaler = StandardScaler()
    x = X_scaler.fit_transform(x)
    pca = PCA(n_components=0.9)
    pca.fit(x)
    pca.transform(x)
    li = list(map(lambda x: sum(x)/len(x), list(pca.transform(x))))
    return li

def build_X_y(reviews, vectors):
    X = []
    Y = []

    for review, polarity in reviews:
        x = convert_text_to_vector(review, vectors)
        y = 0 if polarity == 'negative' else 1
        X.append(x)
        Y.append(y)

    return X, Y


def train_and_test(X_train, y_train, X_test, y_test):
    classifier = LinearSVC()

    classifier.fit(X_train, y_train)
    accuracy = classifier.score(X_test, y_test)
    print(f'accuracy is {accuracy:.4f}')

    return classifier


# 使用50维预训练的词向量
for str in['50d','100d','200d','300d']:
    # glove_input_file = 'D:\A\glove_vector\wiki\glove.6B.'+str+'.txt'
    word2vec_output_file = 'D:\A\glove_vector\wiki\glove.6B.'+str+'.word2vec.txt'
    # glove2word2vec(glove_input_file, word2vec_output_file)

    glove_vectors = KeyedVectors.load_word2vec_format(word2vec_output_file, binary=False)
    # print(glove_vectors.most_similar('twitter'))

    train_reviews, test_reviews = load_movie_reviews()
    print('train:', len(train_reviews))
    print('test:', len(test_reviews))

    X_train, y_train = build_X_y(train_reviews, glove_vectors)
    X_test, y_test = build_X_y(test_reviews, glove_vectors)

    train_and_test(X_train, y_train, X_test, y_test)